import os
import numpy as np
import keras
from keras import layers
from tensorflow import data as tf_data
import matplotlib.pyplot as plt

# Path to the dataset
DOG_CAT_DATASET_PATH = 'C:/Workspace/datasets/kagglecatsanddogs_5340/PetImages'


def del_corrupted_images():
    """
    Delete corrupted images
    :return: None
    """
    num_skipped = 0
    for folder_name in ("Cat", "Dog"):
        folder_path = os.path.join(DOG_CAT_DATASET_PATH, folder_name)
        for fname in os.listdir(folder_path):
            fpath = os.path.join(folder_path, fname)
            try:
                fobj = open(fpath, "rb")
                is_jfif = b"JFIF" in fobj.peek(10)
            finally:
                fobj.close()

            if not is_jfif:
                num_skipped += 1
                # Delete corrupted image
                os.remove(fpath)
    print(f"Deleted {num_skipped} images.")


def create_learning_dataset():
    """
    Create a learning dataset
    :return train_ds, val_ds
    """
    image_size = (180, 180)
    batch_size = 128

    return keras.utils.image_dataset_from_directory(
        DOG_CAT_DATASET_PATH,
        validation_split=0.2,
        subset="both",
        seed=1337,
        image_size=image_size,
        batch_size=batch_size,
    )


data_augmentation_layers = [
    layers.RandomFlip("horizontal"),
    layers.RandomRotation(0.1),
]


def data_augmentation(images):
    for layer in data_augmentation_layers:
        images = layer(images)
    return images


if __name__ == '__main__':
    # Delete corrupted images
    del_corrupted_images()
    # Create a learning dataset
    train_ds, val_ds = create_learning_dataset()
    plt.figure(figsize=(10, 10))
    for images, labels in train_ds.take(1):
        for i in range(9):
            ax = plt.subplot(3, 3, i + 1)
            plt.imshow(np.array(images[i]).astype("uint8"))
            plt.title(int(labels[i]))
            plt.axis("off")

    plt.figure(figsize=(10, 10))
    for images, _ in train_ds.take(1):
        for i in range(9):
            augmented_images = data_augmentation(images)
            ax = plt.subplot(3, 3, i + 1)
            plt.imshow(np.array(augmented_images[0]).astype("uint8"))
            plt.axis("off")
